"""
通过cnn实战mnist
"""
import torch
from torch import nn, optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch.nn.functional as F
from PIL import Image
import matplotlib.pyplot as plt

# 1. 数据准备：
# 1.1数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),  # 转为Tensor，并归一化到[0,1]
    transforms.Normalize((0.5,), (0.5,))  # 标准化到[-1,1]（单通道）
])
# 1.2 数据集
train_dataset = datasets.MNIST(root='../../data',
                               train=True,
                               transform=transforms.ToTensor(),
                               download=False)
test_dataset = datasets.MNIST(root='../../data',
                              train=False,
                              transform=transforms.ToTensor(),
                              download=False)
# 1.3 数据加载器
train_loader = DataLoader(dataset=train_dataset,
                          batch_size=32,
                          shuffle=True)
test_loader = DataLoader(dataset=test_dataset,
                         batch_size=32,
                         shuffle=False)


# 2. 构建模型
# 2.1 模型定义
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 卷积层1：输入1通道（灰度图），输出16通道，3x3卷积核
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1)
        # 池化层：2x2窗口，步长2（尺寸减半）
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        # 卷积层2：输入16通道，输出32通道
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        # 全连接层
        self.fc1 = nn.Linear(32 * 7 * 7, 128)  # MNIST两次池化后为7x7
        self.fc2 = nn.Linear(128, 10)  # 输出10分类

    def forward(self, x):
        # 卷积 → ReLU → 池化
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        # 展平
        x = x.view(-1, 32 * 7 * 7)
        # 全连接层
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


# 2.2 模型定义
cnn = CNN()

# 3. 训练模型
# 3.1 损失函数
criterion = nn.CrossEntropyLoss()
# 3.2 优化器
optimizer = optim.Adam(cnn.parameters(), lr=0.001)


# 3.2 训练循环
def model_train():
    cnn.train()
    for epoch in range(10):
        print(f"第 {epoch + 1} 轮训练开始...")
        train_loss_epoch = 0.0  # 记录整轮的总损失（所有batch的损失之和）

        for i, (inputs, labels) in enumerate(train_loader):
            # 梯度清零
            optimizer.zero_grad()
            # 前向传播
            outputs = cnn(inputs)
            # 计算损失
            loss = criterion(outputs, labels)
            # 反向传播
            loss.backward()
            # 更新参数
            optimizer.step()

            # 累加损失（用于计算整轮平均损失）
            train_loss_epoch += loss.item()

            # 打印训练信息（每100个batch打印一次）
            if (i + 1) % 100 == 0:
                avg_loss_i = train_loss_epoch / (i + 1)  # 计算当前平均损失（可选）
                print(f"第 {epoch + 1} 轮，第 {i + 1} 个batch，当前平均损失：{avg_loss_i:.4f}")

        # 打印整轮平均损失
        avg_loss_epoch = train_loss_epoch / len(train_loader)  # 批次平均损失
        print(f"第 {epoch + 1} 轮训练结束，平均损失：{avg_loss_epoch:.4f}")


# 测试模型
def model_test():
    cnn.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in test_loader:
            outputs = cnn(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    accuracy = 100 * correct / total
    print(f"测试集准确率：{accuracy:.2f}%")


# 模型预测函数
def predict_image(model, img_path):
    # 图像预处理
    img = Image.open(img_path).resize((28, 28)).convert('L')
    img_tensor = transforms.ToTensor()(img)

    # 自动颜色校正（MNIST是黑底白字）
    if img_tensor.mean() > 0.5:
        img_tensor = 1 - img_tensor

    # 添加batch维度并预测
    img_tensor = img_tensor.unsqueeze(0)
    model.eval()
    with torch.no_grad():
        output = model(img_tensor)
        probs = torch.softmax(output, 1)[0]
        pred = torch.argmax(probs).item()

    # 输出结果
    print("\n预测结果详情:")
    print(f"预测数字: {pred}")
    print("各类别概率:")
    for i, p in enumerate(probs):
        print(f"数字{i}: {p.item():.2%}")

    # 可视化
    plt.imshow(img_tensor.squeeze(), cmap='gray')
    plt.title(f"Predicted: {pred}")
    plt.show()
    return pred


if __name__ == '__main__':
    # model_train()
    # model_test()  # 测试集准确率：99.14%

    # 保存模型
    model_save_path = '../model_save/4_2_cnn_mnist.pth'
    # torch.save(cnn.state_dict(), model_save_path)
    # print(f"模型保存成功: {model_save_path}")

    # 预测图像
    predict_image(cnn, '../../images/9.jpg')
